(* # =================================================================== 
   # Vector Type
   #
   # Introduction
   #
   # We build a vector type based on pair.
   #
   # Two basic function: mk_vec v_idx
   #
   #   mk_vec : create a vector with a function from index to element
   #
   #   v_idx : access a element of array
   #
   # We define other functions on vector with the two functions above.
   #
   #   v_map v_reduce v_zip v_split v_join v_unzip1 v_unzip2
   #
   # There are some properties between these functions.
   #
   #  mk_vec_bij_fin: forall (m:nat) (x0 : A) (v : vector m),
   #    mk_vec (v_nth x0 v) = v.
   #
   #  v_nth_mk_vec: forall (n:nat) (f:fin n -> A) (x0:A) (i:fin n),
   #    i < n -> v_nth x0 (mk_vec f) i = f (nat2fin n i).
   #
   #  v_nth_map: forall (A B:Set) (n:nat) (f: A->B) (v:vector A n) (x0:A) (x1:B)(i:fin n),
   #    i < n -> v_nth x1 (v_map x0 f v) i = f (v_nth x0 v i).
   #
   #  v_map_zip : forall n (A B C:Set) (x0:A) (x1:B) (f:A->B->C) (v1:vector A n) (v2:vector B n),
   #    v_map (x0,x1) (cur f) (v_zip x0 x1 v1 v2) = mk_vec (fun i => f (v_nth x0 v1 i) (v_nth x1 v2 i)).
   #
   #  v_map_fg : forall n (A B C:Set) (x0:A) (x1:B) (f:A->B) (g:B->C) (v1:vector A n),
   #    v_map x1 g (v_map x0 f v1) = v_map x0 (fun a => g (f a)) v1.
   #
   #  v_map_map_split: forall (m n : nat) (A B:Set) (x0:vector A n) (x1:A) (f : A -> B) (v : vector A (m*n)),
   #    v_map x0 (fun v0 => v_map x1 f v0) (v_split x1 v) =
   #    mk_vec (fun i => mk_vec (fun j => f (v_idx x1 v (nat2fin (m*n) (i * n + j))))).
   #
   #  v_map_join : forall (m n : nat) (A B:Set) (x0:A) (x1: vector A n) (f : A -> B) (v : vector (vector A n) m),
   #    v_map x0 f (v_join x0 x1 v) 
   #    = mk_vec (fun i => f (v_idx x0 (v_idx x1 v (nat2fin m (i/n))) (nat2fin n (i mod n)))).
   #
   #  v_map_join_split: forall [m n:nat] (A B:Set) (f:A -> B) (x0:A)(x1:vector A n) (x2: B) (x3: vector B n)(v:vector A (m*n)),
   #    v_map x0 f v = v_join x2 x3 (v_map x1 (v_map x0 f) (v_split x0 v)).
   #
   #  v_split_join: forall (A:Set) (m n:nat) (x0:A) (x1:vector A n) (v:vector (vector A n) m),
   #    v_split x0 (v_join x0 x1 v) = v.
   #
   #  v_zip_map: forall (A B:Set) (n:nat) (f g: A -> B) (x0:A) (x1:B) (a b: vector A n),
   #    v_zip x1 x1 (v_map x0 f a) (v_map x0 g b) = v_map (x0,x0) (fun xy => (f (fst xy),g (snd xy))) (v_zip x0 x0 a b).
   #
   #  v_unzip1_mk_vec: forall (A B:Set) n (x0:A*B) f g,
   #    v_unzip1 x0 (mk_vec (fun i => (f i, g i))) = mk_vec (fun i:fin n => f i).
   #
   #  v_unzip2_mk_vec: forall (A B:Set) n (x0:A*B)f g,
   #    v_unzip2 x0 (mk_vec (fun i => (f i, g i))) = mk_vec (fun i: fin n => g i).
   #   
   # =================================================================== *)

Require Export Logic.FunctionalExtensionality.
Require Import List Reals String.
Import ListNotations.
Require Export finType.
Set Implicit Arguments.

Ltac fun_ext:= apply functional_extensionality; intros.

Notation "[ x1 , .. , xn ]" := (.. (tt,x1) .., xn). 

Section VectorType.

Variable A:Set.

(* definition of vector *)
Fixpoint vector (n : nat) : Set :=
  match n with
  | O => unit
  | S n => ((vector n) * A)%type
  end.

Section Vfun.

Fixpoint mk_vec {m:nat} : (fin m -> A) -> vector m:=
  match m with
  | 0 => fun ( _:_ ) => tt :vector 0
  | S m' => fun ( u: fin (S m') -> A) =>
      (mk_vec (fun i => u (fin2finS i)),
       u (exist _ m' (Nat.lt_succ_diag_r m')))
  end.

Lemma mk_vec_Sn : forall (m:nat) (u : fin (S m) -> A),
  mk_vec u = (mk_vec (fun i => u (fin2finS i)),
  u (exist _ m (Nat.lt_succ_diag_r m))).
Proof. auto. Qed.

Fixpoint v_nth {n : nat} (x0 : A) : (vector n) -> nat -> A 
  := match n with
  | O => fun (_ : vector O) (_ : nat) => x0
  | S n' => fun (v : vector (S n')) (i : nat) => 
     if i =? n' then snd v else v_nth x0 (fst v) i
  end.

Definition v_idx {n: nat} (x0 : A) : (vector n) -> fin n -> A :=
  fun v i => v_nth x0 v i.

Lemma v_nth_Sn : forall m (a: vector m) b x x0,
  x < m -> @v_nth (S m) x0 (a,b) x = v_nth x0 a x.
Proof.
  simpl. intros. replace (x=?m) with false. auto. symmetry.
  apply Nat.eqb_neq. lia.
Qed.

Lemma v_nth_x : forall x (a: vector x) b x0,
  @v_nth (S x) x0 (a,b) x = b.
Proof.
  simpl. intros.
  rewrite Nat.eqb_refl. auto.
Qed.

Lemma v_nth_x_fin : forall x (a: vector x) b x0 (i: fin (S x)),
  fin2nat i = x -> @v_nth (S x) x0 (a,b) x = b.
Proof.
  simpl. intros.
  rewrite Nat.eqb_refl. auto.
Qed.

Lemma mk_vec_bij_fin: forall (m:nat) (x0 : A) (v : vector m),
  mk_vec (v_nth x0 v) = v.
Proof. 
  induction m,v. + auto. + simpl;intros. rewrite Nat.eqb_refl.
  f_equal. rewrite <- IHm with(x0:=x0).
  induction m. auto. f_equal. fun_ext. 
  induction x. simpl. replace (x =? S m) with false. auto. symmetry.
  apply Nat.eqb_neq. lia.
Qed.

Lemma v_nth_mk_vec: forall (n:nat) (f:fin n -> A) (x0:A) (i:fin n),
  i < n -> v_nth x0 (mk_vec f) i = f (nat2fin n i).
Proof.
  induction n. lia.
  induction i. autorewrite with core. intros.
  rewrite mk_vec_Sn. destruct (x=?n) eqn:E.
  rewrite <- (v_nth_x n (mk_vec (fun i : fin n => f (fin2finS i))) (f (nat2fin (S n) x)) x0 ) .
  f_equal. f_equal. apply Nat.eqb_eq in E. subst. f_equal.
  apply fin_eq with (n:=S n)(i1:=exist (fun i : nat => i < S n) n (Nat.lt_succ_diag_r n)).
  autorewrite with core. auto. lia. apply Nat.eqb_eq. auto. apply Nat.eqb_neq in E.
  rewrite v_nth_Sn. rewrite <- fin2nat_nat2fin with (n:= n)(i:=x).
  rewrite IHn with (f:=fun i : fin n => f (fin2finS i)).
  f_equal. apply fin_eq. autorewrite with core. auto.
  auto. lia. autorewrite with core; lia; lia. autorewrite with core; lia; lia.
  lia. lia.
Qed.

End Vfun.
End VectorType.

Arguments mk_vec {A} {m}.
Arguments v_nth {A} {n}.

Hint Rewrite mk_vec_Sn.

Section VecFun.

Definition cur {A B C:Set} (f: A -> B -> C) :=
  fun xy : A * B => f (fst xy) (snd xy).

Definition uncur {A B C:Set} (f: A*B -> C) :=
  fun x y => f (x,y).

Definition v_make {A : Set} (n : nat) (x0:A): vector A n :=
  mk_vec (fun x => x0).

(* mapping one vector to another *)
Definition v_map {A B : Set} {n : nat} (x0:A) (f : A -> B) (v:vector A n) : vector B n :=
  mk_vec (fun x => f (v_idx x0 v x)).

(* fold a vector to an element from left to right *)
Fixpoint v_reduce {A B : Set} {n : nat} (f : A -> B -> B) (init_val:B)
  : (vector A n) -> B :=
  match n with
  | O => fun (v : vector A 0) => init_val
  | S n' => fun (v : vector A (S n')) => 
    f (snd v) (v_reduce f init_val (fst v))
  end.

(* combine two vector into a vector of pairs *)
Definition v_zip {A B:Set} {n:nat} (x0:A) (x1:B) (v1:vector A n) (v2:vector B n) : vector (A*B) n:=
  mk_vec (fun x => (v_idx x0 v1 x,v_nth x1 v2 x)).

Definition v_unzip1 {A B:Set} {n:nat} (x0:A*B) (v: vector (A*B) n): vector A n :=
  mk_vec (fun x => fst (v_idx x0 v x)).

Definition v_unzip2 {A B:Set} {n:nat} (x0:A*B) (v: vector (A*B) n): vector B n :=
  mk_vec (fun x => snd (v_idx x0 v x)).

Hint Unfold v_map v_zip v_unzip1 v_unzip2 cur uncur v_idx.

Lemma v_unzip1_mk_vec: forall (A B:Set) n (x0:A*B) f g,
  v_unzip1 x0 (mk_vec (fun i => (f i, g i))) = 
  mk_vec (fun i:fin n => f i).
Proof.
  autounfold.
  induction n;auto.
  intros. f_equal. fun_ext.
  rewrite v_nth_mk_vec;auto. unfold fst.
  f_equal. induction x;auto.
Qed.

Lemma v_unzip2_mk_vec: forall (A B:Set) n (x0:A*B)f g,
  v_unzip2 x0 (mk_vec (fun i => (f i, g i))) = 
  mk_vec (fun i: fin n => g i).
Proof.
  autounfold.
  induction n;auto.
  intros. f_equal. fun_ext.
  rewrite v_nth_mk_vec;auto. unfold snd.
  f_equal. induction x;auto.
Qed.

Lemma v_nth_map: forall (A B:Set) (n:nat) (f: A->B) (v:vector A n) (x0:A) (x1:B)(i:fin n),
  i < n -> v_nth x1 (v_map x0 f v) i = f (v_nth x0 v i).
Proof.
  unfold v_map,v_idx. intros. rewrite v_nth_mk_vec;auto. repeat f_equal.
  apply fin_eq. autorewrite with core;auto.
Qed.

Lemma v_map_zip : forall n (A B C:Set) (x0:A) (x1:B) (f:A->B->C) (v1:vector A n) (v2:vector B n),
  v_map (x0,x1) (cur f) (v_zip x0 x1 v1 v2) = mk_vec (fun i => f (v_nth x0 v1 i) (v_nth x1 v2 i)).
Proof.
  autounfold. induction n;auto.
  induction v1,v2. rewrite ?mk_vec_Sn. f_equal.
  induction n;auto.
  replace (mk_vec (fun i : fin (S n) => f (v_nth x0 ((a, b):vector A (S (S n))) (fin2finS i)) 
    (v_nth x1 ((v, b0): vector B (S (S n))) (fin2finS i)))) with
    (mk_vec (fun i : fin (S n) => f (v_nth x0 a i) (v_nth x1 v i))).
  rewrite <- IHn. f_equal. fun_ext. rewrite fin2nat_exist. rewrite v_nth_Sn.
  replace (fun i : fin (S n) => (v_nth x0 ((a, b):vector A (S (S n))) (fin2finS i), 
    v_nth x1 ((v, b0):vector B (S (S n))) (fin2finS i)))
    with (fun x2 : fin (S n) => (v_nth x0 a x2, v_nth x1 v x2)).
  auto. fun_ext. 
  f_equal;rewrite v_nth_Sn;autorewrite with core; auto;
  autorewrite with core; auto.
  autorewrite with core; auto.
  f_equal. fun_ext. f_equal; rewrite v_nth_Sn;autorewrite with core; auto;
  autorewrite with core; auto. rewrite fin2nat_exist. rewrite v_nth_x.
  auto.
Qed.

Lemma mk_vec_zip : forall n (A B C:Set) (x0:A) (x1:B) (f:A->B->C) (g: fin n-> fin n) (v1:vector A n) (v2:vector B n),
  mk_vec (fun i => cur f (v_nth (x0,x1) (v_zip x0 x1 v1 v2) (g i))) 
  = mk_vec (fun i => f (v_nth x0 v1 (g i))(v_nth x1 v2 (g i))).
Proof. autounfold.
  induction n;auto. intros. f_equal. fun_ext. f_equal. 
  rewrite v_nth_mk_vec. unfold fst. f_equal;rewrite fin2nat_nat2fin;auto. auto.
  rewrite v_nth_mk_vec. unfold snd. f_equal;rewrite fin2nat_nat2fin;auto. auto.
Qed. 

Lemma mk_vec2_zip : forall m n (A B C:Set) (x0:A) (x1:B) (f:A->B->C) (g: fin m -> fin n -> fin (m*n)) (v1:vector A (m*n)) (v2:vector B (m*n)),
  mk_vec (fun i => mk_vec (fun j => cur f (v_nth (x0,x1) (v_zip x0 x1 v1 v2) (g i j)))) 
  = mk_vec (fun i => mk_vec (fun j => f (v_nth x0 v1 (g i j))(v_nth x1 v2 (g i j)))).
Proof. autounfold. induction m;auto. intros. f_equal. fun_ext.
  induction n;auto. intros. f_equal. fun_ext.
  rewrite v_nth_mk_vec. f_equal;rewrite fin2nat_nat2fin;auto. induction (g x x2).
  simpl. auto. induction (g x x2); simpl;auto.
  induction (g x x2); simpl; auto.
Qed.

Lemma v_map_fg : forall n (A B C:Set) (x0:A) (x1:B) (f:A->B) (g:B->C) (v1:vector A n),
  v_map x1 g (v_map x0 f v1) = v_map x0 (fun a => g (f a)) v1.
Proof.
  autounfold. induction n;auto. intros. f_equal. fun_ext.
  f_equal. rewrite v_nth_mk_vec;auto. f_equal. rewrite fin2nat_nat2fin;auto.
Qed.

Definition v_split {A : Set} {m n:nat} (x0 : A) (v:vector A (m*n)) : vector (vector A n) m :=
  mk_vec (fun i => mk_vec (fun j => v_idx x0 v (nat2fin (m*n) (i*n+j)))).

Definition v_join {A : Set} {m n:nat} (x0:A) (x1:vector A n) (v:vector (vector A n) m) : vector A (m * n):=
  mk_vec (fun x => v_idx x0 (v_nth x1 v (nat2fin m (x/n))) (nat2fin n (x mod n))).

Hint Unfold v_split v_join.

Lemma v_map_map_split: forall (m n : nat) (A B:Set) (x0:vector A n) (x1:A) (f : A -> B) (v : vector A (m*n)),
  v_map x0 (fun v0 => v_map x1 f v0) (v_split x1 v) =
  mk_vec (fun i => mk_vec (fun j => f (v_idx x1 v (nat2fin (m*n) (i * n + j))))).
Proof. autounfold.
  intros. 
  induction m;auto. f_equal. fun_ext. 
  induction n;auto. f_equal. fun_ext.
  f_equal.  intros. rewrite ?v_nth_mk_vec;auto.
  f_equal. 
  assert (x * S n + x2 < S m * S n).
  induction x,x2;autorewrite with core; auto.
  assert (x * S n + x2 <= m * S n + n).
  apply plus_le_compat. apply Nat.mul_le_mono_r. lia. lia.
  apply Nat.le_lt_trans with (m:=m * S n + n). auto.
  simpl. rewrite <- Nat.add_succ_l. rewrite (Nat.add_comm (S n) (m*S n)) .
  apply plus_lt_compat_l. lia. 
  autorewrite with core;
  induction x,x2;autorewrite with core; auto; try lia.
Qed.

Lemma v_map_join : forall (m n : nat) (A B:Set) (x0:A) (x1: vector A n) (f : A -> B) (v : vector (vector A n) m),
  v_map x0 f (v_join x0 x1 v) 
  = mk_vec (fun i => f (v_idx x0 (v_idx x1 v (nat2fin m (i/n))) (nat2fin n (i mod n)))).
Proof.
  autounfold.
  intros. induction m;auto.
  induction n. f_equal. rewrite Nat.mul_0_r. simpl. auto.
  f_equal. fun_ext. f_equal. rewrite v_nth_mk_vec;induction x;auto. 
  repeat f_equal;autorewrite with core;induction x;auto.
Qed.

Lemma v_map_join_split: forall [m n:nat] (A B:Set) (f:A -> B) 
  (x0:A)(x1:vector A n) (x2: B) (x3: vector B n)(v:vector A (m*n)),
  v_map x0 f v = v_join x2 x3 (v_map x1 (v_map x0 f) (v_split x0 v)).
Proof.
  autounfold.
  induction m;auto. induction n. simpl. rewrite Nat.mul_0_r. auto.
  intros. f_equal. fun_ext. 
  rewrite ?v_nth_mk_vec;auto. repeat f_equal.
  assert (x / S n * S n + x mod S n = x).
  rewrite Nat.mul_comm. rewrite Nat.div_mod with (y:= S n);auto.
  assert (x < S m * S n). induction x;auto.
  autorewrite with core;auto. rewrite H. apply fin_eq. autorewrite with core;auto.
  apply Nat.mod_upper_bound;auto. apply Nat.div_lt_upper_bound;auto.
  rewrite Nat.mul_comm. auto.
Qed.

Lemma v_reduce_map: forall [n:nat] (A B C:Set) (g:A -> B) 
    (f: B -> C -> C) (i:C) (x0:A) (v:vector A n) ,
  v_reduce f i (v_map x0 g v) = v_reduce (fun x a => f (g x) a) i v.
Proof.
  autounfold. induction n,v;auto. simpl. f_equal. rewrite Nat.eqb_refl. auto.
  rewrite <-IHn with (x0:=x0) (v:=v). f_equal. induction n;auto.
  f_equal. fun_ext. f_equal. assert (x< S n);auto. replace (x =? S n)%nat with false.
  auto. symmetry. rewrite Nat.eqb_neq. lia.
Qed.

Lemma v_join_split: forall (A:Set) (m n:nat) (x0:A) (x1:vector A n) (v:vector A (m*n)),
  v_join x0 x1 (v_split x0 v) = v.
Proof.
  autounfold. induction m. simpl. induction v;auto.
  induction n. simpl. rewrite Nat.mul_0_r. induction v. auto.
  intros. rewrite <- mk_vec_bij_fin with(x0:=x0).
  f_equal. fun_ext. rewrite ?v_nth_mk_vec;auto.
  f_equal. autorewrite with core;auto.
  rewrite Nat.mul_comm. rewrite Nat.div_mod with(y:= S n);auto.
  apply Nat.mod_upper_bound; auto. 
  apply Nat.div_lt_upper_bound;auto. 
  rewrite Nat.mul_comm;induction x;auto.
  autorewrite with core;auto. rewrite Nat.mul_comm.
  rewrite <- Nat.div_mod;auto. induction x;auto.
  apply Nat.mod_upper_bound;auto.
  apply Nat.div_lt_upper_bound;auto. 
  rewrite Nat.mul_comm;induction x;auto.
Qed.

Lemma mk_vec_tt: forall (m:nat) (a: vector unit m),
  mk_vec (fun _ : fin m => tt) = a.
Proof.
  induction m,a;auto.
  simpl. f_equal. apply IHm with (a:=v). induction u;auto.
Qed.

Lemma v_split_join: forall (A:Set) (m n:nat) (x0:A) (x1:vector A n) (v:vector (vector A n) m),
  v_split x0 (v_join x0 x1 v) = v.
Proof.
  autounfold. induction m. simpl. induction v;auto.
  induction n. rewrite Nat.mul_0_r. simpl. induction v;auto.
  rewrite mk_vec_tt with(a:=a). induction b;auto.
  intros. rewrite <- mk_vec_bij_fin with(x0:=x1).
  f_equal. fun_ext. rewrite <- mk_vec_bij_fin with(x0:=x0).
  f_equal. fun_ext.
  assert (x * S n + x2 < S m * S n).
  induction x,x2;autorewrite with core; auto.
  assert (x * S n + x2 <= m * S n + n).
  apply plus_le_compat. apply Nat.mul_le_mono_r. lia. lia.
  apply Nat.le_lt_trans with (m:=m * S n + n). auto.
  simpl. rewrite <- Nat.add_succ_l. rewrite (Nat.add_comm (S n) (m*S n)) .
  apply plus_lt_compat_l. lia. 
  rewrite v_nth_mk_vec;auto.
  autorewrite with core;auto. repeat f_equal.
  rewrite Nat.add_comm,Nat.div_add;auto. rewrite Nat.div_small;auto.
  rewrite Nat.add_comm,Nat.mod_add,Nat.mod_small;auto.
  rewrite Nat.add_comm,Nat.mod_add,Nat.mod_small;auto. 
  autorewrite with core;auto. autorewrite with core;auto.
  rewrite Nat.add_comm,Nat.div_add;auto. rewrite Nat.div_small;auto.
  rewrite Nat.add_0_l. auto.
  autorewrite with core;auto. autorewrite with core;auto.
Qed.

Lemma v_zip_map: forall (A B:Set) (n:nat) (f g: A -> B) (x0:A) (x1:B) (a b: vector A n),
  v_zip x1 x1 (v_map x0 f a) (v_map x0 g b) = v_map (x0,x0) (fun xy => (f (fst xy),g (snd xy))) (v_zip x0 x0 a b).
Proof.
  autounfold. induction n;auto. intros. f_equal.
  fun_ext. rewrite ?v_nth_mk_vec;auto.
Qed.

End VecFun.



